AllGather ================= 将各输入分片的数据进行 **全量收集(All-Gather)**, 并在输出中按指定维度进行拼接复制。 该算子通常用于并行/多核场景中,将每个 rank 的本地数据 收集成一个完整的输出张量。 算子行为可概括为两步: 1. 将 ``input`` 中每个 ``input_rank`` 对应的数据块拷贝到 ``output`` 的首段连续区域; 2. 将首段完整数据复制到其余 ``output_rank - 1`` 个位置,形成完整聚合结果。 输入: - **input** - 输入数据地址,按 ``[input_rank, data_size]`` 形式连续存储。 - **input_rank** - 输入 rank 数量(通常对应参与 AllGather 的并行单元数)。 - **output_rank** - 输出 rank 数量(通常等于或大于 ``input_rank``)。 - **data_size** - 单个 rank 对应的数据元素个数。 - **core_mask** - 核掩码(仅适用于共享存储版本)。 输出: - **output** - AllGather 结果输出地址, 数据按 ``[output_rank, input_rank, data_size]`` 的逻辑顺序存储。 支持平台: ``FT78NE`` ``MT7004`` .. note:: - FT78NE 支持 ``fp32``、``fp64``、``int8``、``int16``、``int32``、``cplx64``、``cplx128`` 类型 - MT7004 支持 ``fp16``、``fp32``、``int16``、``int32``、``cplx64`` 类型 - 当前实现基于 DMA进行数据搬运 - 输入与输出地址需保证足够的连续空间以容纳聚合结果 **共享存储版本:** .. c:function:: void fp_allgather_s(float* input, float* output, int input_rank, int output_rank, int data_size, int core_mask) .. c:function:: void dp_allgather_s(double* input, double* output, int input_rank, int output_rank, int data_size, int core_mask) .. c:function:: void i8_allgather_s(int8_t* input, int8_t* output, int input_rank, int output_rank, int data_size, int core_mask) .. c:function:: void i16_allgather_s(int16_t* input, int16_t* output, int input_rank, int output_rank, int data_size, int core_mask) .. c:function:: void i32_allgather_s(int32_t* input, int32_t* output, int input_rank, int output_rank, int data_size, int core_mask) .. c:function:: void c64_allgather_s(cplx64* input, cplx64* output, int input_rank, int output_rank, int data_size, int core_mask) .. c:function:: void c128_allgather_s(cplx128* input, cplx128* output, int input_rank, int output_rank, int data_size, int core_mask) **C调用示例:** .. code-block:: c :linenos: :emphasize-lines: 14-16 // FT78NE 示例 #include #include int main(int argc, char* argv[]) { float *input = (float *)0xA0000000; float *output = (float *)0xC0000000; int input_rank = 4; int output_rank = 4; int data_size = 256; int core_mask = 0xff; fp_allgather_s(input, output, input_rank, output_rank, data_size, core_mask); return 0; } **私有存储版本:** .. c:function:: void fp_allgather_p(float* input, float* output, int input_rank, int output_rank, int data_size) .. c:function:: void dp_allgather_p(double* input, double* output, int input_rank, int output_rank, int data_size) .. c:function:: void i8_allgather_p(int8_t* input, int8_t* output, int input_rank, int output_rank, int data_size) .. c:function:: void i16_allgather_p(int16_t* input, int16_t* output, int input_rank, int output_rank, int data_size) .. c:function:: void i32_allgather_p(int32_t* input, int32_t* output, int input_rank, int output_rank, int data_size) .. c:function:: void c64_allgather_p(cplx64* input, cplx64* output, int input_rank, int output_rank, int data_size) .. c:function:: void c128_allgather_p(cplx128* input, cplx128* output, int input_rank, int output_rank, int data_size) **C调用示例:** .. code-block:: c :linenos: :emphasize-lines: 13-15 // FT78NE 示例 #include #include int main(int argc, char* argv[]) { float *input = (float *)0x10810000; // L2 空间 float *output = (float *)0x10820000; int input_rank = 4; int output_rank = 4; int data_size = 256; fp_allgather_p(input, output, input_rank, output_rank, data_size); return 0; }